package io.gitee.milklynk;

import io.gitee.milklynk.vo.response.ChatResponse;
import io.gitee.milklynk.vo.response.Choice;
import io.gitee.milklynk.vo.response.Usage;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

public class ChatResponseMerger {

    public static ChatResponse merge(List<ChatResponse> responses) {
        ChatResponse merged = new ChatResponse();
        if (responses == null || responses.isEmpty()) {
            return merged;
        }

        for (ChatResponse response : responses) {
            // 合并基础字段
            if (merged.getId() == null) {
                merged.setId(response.getId());
            }
            if (merged.getCreated() == null) {
                merged.setCreated(response.getCreated());
            }
            if (merged.getModel() == null) {
                merged.setModel(response.getModel());
            }
            if (merged.getSystemFingerprint() == null) {
                merged.setSystemFingerprint(response.getSystemFingerprint());
            }
            if (merged.getObject() == null) {
                merged.setObject(response.getObject());
            }
            // 假设最后一个usage有效
            if (response.getUsage() != null) {
                merged.setUsage(copyUsage(response.getUsage()));
            }
            // 合并choices
            if (response.getChoices() != null) {
                for (Choice currentChoice : response.getChoices()) {
                    mergeChoice(merged, currentChoice);
                }
            }
        }

        return merged;
    }

    private static void mergeChoice(ChatResponse merged, Choice currentChoice) {
        List<Choice> mergedChoices = merged.getChoices() != null ? merged.getChoices() : new ArrayList<>();
        merged.setChoices(mergedChoices);

        Integer index = currentChoice.getIndex() != null ? currentChoice.getIndex() : 0;
        Choice targetChoice = findOrCreateMergedChoice(mergedChoices, index);

        // 合并finish_reason
        if (currentChoice.getFinishReason() != null) {
            targetChoice.setFinishReason(currentChoice.getFinishReason());
        }

        // 合并message
        if (currentChoice.getMessage() != null) {
            if (targetChoice.getMessage() == null) {
                targetChoice.setMessage(new Choice.Message());
            }
            mergeMessage(targetChoice.getMessage(), currentChoice.getMessage());
        }
    }

    private static Choice findOrCreateMergedChoice(List<Choice> mergedChoices, Integer index) {
        for (Choice choice : mergedChoices) {
            if (Objects.equals(choice.getIndex(), index)) {
                return choice;
            }
        }
        Choice newChoice = new Choice();
        newChoice.setIndex(index);
        mergedChoices.add(newChoice);
        return newChoice;
    }

    private static void mergeMessage(Choice.Message mergedMsg, Choice.Message currentMsg) {
        // Role只设置一次
        if (mergedMsg.getRole() == null) {
            mergedMsg.setRole(currentMsg.getRole());
        }

        // 拼接content
        mergedMsg.setContent(append(mergedMsg.getContent(), currentMsg.getContent()));

        // 拼接reasoning_content
        mergedMsg.setReasoningContent(append(mergedMsg.getReasoningContent(), currentMsg.getReasoningContent()));

        // 合并tool_calls
        if (currentMsg.getToolCalls() != null) {
            if (mergedMsg.getToolCalls() == null) {
                mergedMsg.setToolCalls(new ArrayList<>());
            }
            mergeToolCalls(mergedMsg.getToolCalls(), currentMsg.getToolCalls());
        }
    }

    private static void mergeToolCalls(List<Choice.Message.ToolCall> mergedCalls, List<Choice.Message.ToolCall> currentCalls) {
        for (Choice.Message.ToolCall currentCall : currentCalls) {
            Choice.Message.ToolCall targetCall = findOrCreateToolCall(mergedCalls, currentCall.getId());
            mergeSingleToolCall(targetCall, currentCall);
        }
    }

    private static Choice.Message.ToolCall findOrCreateToolCall(List<Choice.Message.ToolCall> mergedCalls, String id) {
        for (Choice.Message.ToolCall call : mergedCalls) {
            if (Objects.equals(call.getId(), id)) {
                return call;
            }
        }
        Choice.Message.ToolCall newCall = new Choice.Message.ToolCall();
        newCall.setId(id);
        mergedCalls.add(newCall);
        return newCall;
    }

    private static void mergeSingleToolCall(Choice.Message.ToolCall merged, Choice.Message.ToolCall current) {
        // Type只设置一次
        if (merged.getType() == null) {
            merged.setType(current.getType());
        }

        // 合并function
        if (current.getFunction() != null) {
            if (merged.getFunction() == null) {
                merged.setFunction(new Choice.Message.ToolCall.Function());
            }
            mergeFunction(merged.getFunction(), current.getFunction());
        }
    }

    private static void mergeFunction(Choice.Message.ToolCall.Function mergedFunc, Choice.Message.ToolCall.Function currentFunc) {
        // Name只设置一次
        if (mergedFunc.getName() == null) {
            mergedFunc.setName(currentFunc.getName());
        }

        // 拼接arguments
        mergedFunc.setArguments(append(mergedFunc.getArguments(), currentFunc.getArguments()));
    }

    private static String append(String existing, String current) {
        return existing == null ? current : existing + current;
    }

    private static Usage copyUsage(Usage source) {
        if (source == null) return null;
        Usage usage = new Usage();
        usage.setCompletionTokens(source.getCompletionTokens());
        usage.setPromptTokens(source.getPromptTokens());
        usage.setPromptCacheHitTokens(source.getPromptCacheHitTokens());
        usage.setPromptCacheMissTokens(source.getPromptCacheMissTokens());
        usage.setTotalTokens(source.getTotalTokens());
        if (source.getCompletionTokensDetails() != null) {
            Usage.CompletionTokensDetails details = new Usage.CompletionTokensDetails();
            details.setReasoningTokens(source.getCompletionTokensDetails().getReasoningTokens());
            usage.setCompletionTokensDetails(details);
        }
        return usage;
    }
}
